package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"math/rand"
	"os"
)

// =================================================================

/*
FileEncryptor
	定义一个用来进行文件加解密的接口
	下定义这个概念专门用来处理文件写入加密和文件读取解密
*/

type FileEncryptorInterface interface {
	EncryptAndWriteToFile(filename string, plaintext []byte) error
	ReadAndDecryptFile(filename string) ([]byte, error)
}

// 只把分组模式，填充模式和加密算法统一起来

type BlockModes int

// 分组模式
const (
	CBC BlockModes = iota
	CFB
	GCM
)

type Padding int

// 填充模式
const (
	PKCS7 Padding = iota
	PKCS5
	ZeroPadding
)

type Encryption int

const (
	AES Encryption = iota
	ChaCha20
	RSA
	ECC
)

// ========================================================================

// EncryptorList 加密器集合
type EncryptorList struct {
	content []FileEncryptorInterface
}

// Encryptor 具体加密器自定义：AES-256+ CBC 模式（需要 IV）+ PKCS7（因为 AES 只能处理 16 字节的倍数）
type Encryptor struct {
	Key  []byte // 密钥 32 字节
	IV   []byte
	Name string
	Encryption
	Padding
	BlockModes
}

func NewEncryptor(key, iv []byte, name string, encryption Encryption, padding Padding, modes BlockModes) Encryptor {
	// 都先不做过滤
	return Encryptor{
		Key:        key,
		IV:         iv,
		Name:       name,
		Encryption: encryption,
		Padding:    padding,
		BlockModes: modes,
	}
}

// =======================================================================

// Pkcs7Pad :PKCS7 填充，保证数据是 16 字节的倍数
func (e *Encryptor) Pkcs7Pad(data []byte, blockSize int) []byte {
	needPadding := blockSize - (len(data) % blockSize)
	padText := bytes.Repeat([]byte{byte(needPadding)}, needPadding)
	return append(data, padText...)
}

// Pkcs7UnPad 去除填充
func (e *Encryptor) Pkcs7UnPad(data []byte) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, errors.New("invalid padding size")
	}
	padding := int(data[length-1])
	if padding > length {
		return nil, errors.New("invalid padding content")
	}
	return data[:length-padding], nil
}

// EncryptAES256CBC : AES-256 CBC 加密
func (e *Encryptor) EncryptAES256CBC(plaintext []byte) ([]byte, error) {
	block, err := aes.NewCipher(e.Key)
	if err != nil {
		return nil, err
	}

	plaintext = e.Pkcs7Pad(plaintext, aes.BlockSize)

	ciphertext := make([]byte, len(plaintext))
	mode := cipher.NewCBCEncrypter(block, e.IV)
	mode.CryptBlocks(ciphertext, plaintext)

	return ciphertext, nil
}

// DecryptAES256CBC :AES-256 CBC 解密
func (e *Encryptor) DecryptAES256CBC(ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(e.Key)
	if err != nil {
		return nil, err
	}

	if len(ciphertext)%aes.BlockSize != 0 {
		return nil, errors.New("ciphertext is not a multiple of the block size")
	}

	plaintext := make([]byte, len(ciphertext))
	mode := cipher.NewCBCDecrypter(block, e.IV)
	mode.CryptBlocks(plaintext, ciphertext)

	return e.Pkcs7UnPad(plaintext)
}

// EncryptAndWriteToFile 将加密数据写入文件
func (e *Encryptor) EncryptAndWriteToFile(filename string, plaintext []byte) error {
	// 加密数据
	ciphertext, err := e.EncryptAES256CBC(plaintext)
	if err != nil {
		return err
	}

	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write(ciphertext)
	return err
}

// ReadAndDecryptFile 读取加密文件并解密
func (e *Encryptor) ReadAndDecryptFile(filename string) ([]byte, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	// 读取密文
	ciphertext, err := io.ReadAll(file)
	if err != nil {
		return nil, err
	}
	// 解密
	return e.DecryptAES256CBC(ciphertext)
}

// GenerateKey 生成一个随机的 32 字节密钥（AES-256）
func generateKey() []byte {
	key := make([]byte, 32)
	_, err := rand.Read(key)
	if err != nil {
		panic(err)
	}
	return key
}

// GenerateIV 生成一个随机的 16 字节 IV
func generateIV() []byte {
	iv := make([]byte, aes.BlockSize)
	_, err := rand.Read(iv)
	if err != nil {
		panic(err)
	}
	return iv
}

// ===================================================================================================

func main() {
	// 生成 AES-256 密钥
	key := generateKey()
	iv := generateIV()
	fmt.Println("使用的密钥:", base64.StdEncoding.EncodeToString(key))
	fmt.Println("使用的iv:", base64.StdEncoding.EncodeToString(iv))
	// 创建容器，并将自定义的加密器放进去
	encryptorContainer := EncryptorList{content: make([]FileEncryptorInterface, 0)}
	encryptor := NewEncryptor(key, iv, "AES-256", AES, PKCS7, CBC)
	encryptorContainer.content = append(encryptorContainer.content, &encryptor)

	plaintext := make([]byte, 0)
	// 需要加密的数据
	for i := 0; i < 100; i++ {
		plaintext = append(plaintext, []byte("这是一个需要加密的文件内容！")...)
	}

	// 定义文件名
	filename := "encrypted_data.txt"

	// 加密并写入文件
	err := encryptor.EncryptAndWriteToFile(filename, plaintext)
	if err != nil {
		fmt.Println("加密失败:", err)
		return
	}
	fmt.Println("加密并写入文件成功！")

	// 读取并解密
	decryptedText, err := encryptor.ReadAndDecryptFile(filename)
	if err != nil {
		fmt.Println("解密失败:", err)
		return
	}
	fmt.Println("解密成功，内容:", string(decryptedText))
}
